
import json
import os

import pandas as pd
import tushare as ts

def get_tushare_token():

    current_file_path = os.path.abspath(__file__)

    # 计算主目录路径：向上回退3层目录
    project_root = os.path.abspath(
        os.path.join(current_file_path, '../../..')
    )

    # 组合配置文件完整路径
    json_path = os.path.join(project_root, 'config.json')

    try:
        # 检查文件是否存在
        if not os.path.exists(json_path):
            raise FileNotFoundError(f"配置文件 {json_path} 未找到")

        # 读取并解析 JSON
        with open(json_path, 'r', encoding='utf-8') as f:
            config = json.load(f)

        # 逐层提取字段
        if 'tushare' not in config:
            raise ValueError("JSON 文件中缺少 'tushare' 字段")

        token = config['tushare'].get('token')
        if not token:
            raise ValueError("'tushare' 对象中缺少 'token' 字段")

        return token

    except json.JSONDecodeError:
        raise ValueError("JSON 文件格式错误，请检查内容")


def load_data(stock_code='601398.SH', start_date='20100101', end_date='20241231', ex_data=None):
    """增强版数据加载器，支持手动添加数据行"""
    token = get_tushare_token()
    ts.set_token(token)
    pro = ts.pro_api()

    # 获取日线数据（按日期升序排列）
    df = pro.daily(
        ts_code=stock_code,
        start_date=start_date,
        end_date=end_date,
        fields='trade_date,open,high,low,close,vol'
    ).sort_values('trade_date')

    # 获取复权因子（按日期升序排列）
    adj_df = pro.adj_factor(
        ts_code=stock_code,
        start_date=start_date,
        end_date=end_date
    ).sort_values('trade_date')

    # 合并数据
    df = df.merge(adj_df[['trade_date', 'adj_factor']], on='trade_date', how='left')

    # 处理外部数据 --------------------------------------------------
    if ex_data is not None:
        # 数据格式验证
        required_columns = ['trade_date', 'open', 'high', 'low', 'close', 'vol', 'adj_factor']
        if isinstance(ex_data, pd.DataFrame):
            if not all(col in ex_data.columns for col in required_columns):
                missing = [col for col in required_columns if col not in ex_data.columns]
                raise ValueError(f"外部数据缺少必要列: {missing}")
            ex_df = ex_data.copy()
        elif isinstance(ex_data, dict):
            if not all(col in ex_data for col in required_columns):
                missing = [col for col in required_columns if col not in ex_data]
                raise ValueError(f"字典数据缺少必要键: {missing}")
            ex_df = pd.DataFrame([ex_data])
        else:
            raise TypeError("ex_data 必须是 DataFrame 或 Dict 类型")

        # 日期格式统一化
        ex_df['trade_date'] = pd.to_datetime(ex_df['trade_date']).dt.strftime('%Y%m%d')

        # 合并数据（自动去重）
        df = pd.concat([df, ex_df], axis=0).drop_duplicates('trade_date', keep='last').sort_values('trade_date')

    # 复权计算 ------------------------------------------------------
    latest_adj_factor = df['adj_factor'].iloc[-1]
    for col in ['open', 'high', 'low', 'close']:
        df[col] = df[col] * df['adj_factor'] / latest_adj_factor
    df['vol'] = df['vol'] / df['adj_factor']

    # 格式处理
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df.set_index('trade_date', inplace=True)
    df.sort_index(ascending=True, inplace=True)
    df.rename(columns={'vol': 'volume'}, inplace=True)
    df.drop(columns=['adj_factor'], inplace=True)

    return df


def load_etf(ts_code, start_date, end_date):
    token = get_tushare_token()
    ts.set_token(token)
    pro = ts.pro_api()
    try:
        # ========== 1. 获取规范化的原始数据 ==========
        # 使用通用日线接口（更稳定）
        df = pro.fund_daily(
            ts_code=ts_code,
            start_date=start_date,
            end_date=end_date,
            fields='ts_code,trade_date,open,high,low,close,vol,amount'
        )

        if df.empty:
            print("未获取到原始数据，请检查股票代码和时间范围")
            return None

        # 统一日期格式并排序
        df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
        df = df.sort_values('trade_date', ascending=True)  # 确保升序排列
        df = df.reset_index(drop=True)
        df.rename(columns={'vol': 'volume'}, inplace=True)

        # ========== 2. 获取并验证复权因子 ==========
        # 使用增强版复权因子接口
        adj = pro.adj_factor(
            ts_code=ts_code,
            start_date=start_date,
            end_date=end_date
        ).sort_values('trade_date', ascending=True)

        if adj.empty:
            print(f"警告：{ts_code} 无复权因子数据，使用原始数据")
            return df.set_index('trade_date')

        # 合并前处理日期格式
        adj['trade_date'] = pd.to_datetime(adj['trade_date'], format='%Y%m%d')
        merged = pd.merge_asof(
            df.sort_values('trade_date'),
            adj[['trade_date', 'adj_factor']].sort_values('trade_date'),
            on='trade_date',
            direction='forward'  # 关键：向前匹配最近因子
        )

        # ========== 3. 数据校验与修复 ==========
        # 处理可能的缺失值（使用前向填充）
        merged['adj_factor'] = merged['adj_factor'].ffill().bfill().fillna(1.0)

        # 强制类型转换
        merged['adj_factor'] = merged['adj_factor'].astype(float)
        price_cols = ['open', 'high', 'low', 'close']
        merged[price_cols] = merged[price_cols].apply(pd.to_numeric, errors='coerce')

        # ========== 4. 精确复权计算 ==========
        # 按时间顺序计算复权比例
        merged['adj_ratio'] = merged['adj_factor'].iloc[-1] / merged['adj_factor']

        # 计算复权价格（保留4位小数计算，最终保留2位）
        merged[price_cols] = (merged[price_cols].values * merged['adj_ratio'].values.reshape(-1, 1)).round(4)
        merged[price_cols] = merged[price_cols].round(2)

        # ========== 5. 结果格式化 ==========
        final_df = merged.set_index('trade_date').sort_index(ascending=True)
        final_df = final_df[['ts_code'] + price_cols + ['vol', 'amount', 'adj_factor']]

        print(f"\n处理成功：{ts_code}")
        print(f"复权因子范围：{final_df['adj_factor'].min():.4f} - {final_df['adj_factor'].max():.4f}")
        return final_df

    except Exception as e:
        print(f"处理失败：{str(e)}")
        return None

if __name__ == '__main__':
    load_data()